import os, pickle, json, random, argparse
from gt_visualizer import gt_visualizer, get_box_list
from pred_visualizer import pred_visualizer, nuscenes_info_loader
from fp_filter import fp_filter

def visualize(NUSCENES_PATH, SAMPLE_DATA_PATH, PRED_PATH, SAVE_PATH, sample_num):
    # prepare dir
    if not os.path.exists(SAVE_PATH):
        os.makedirs(SAVE_PATH, exist_ok=True)
    with open(os.path.join(NUSCENES_PATH, 'nuscenes_infos_val.pkl'), 'rb') as f:
        gt_samples = pickle.load(f)
    with open(PRED_PATH, 'r') as f:
        pred_annos = json.load(f)
    sample_annos_dict = pred_annos['results']
    # random sample
    gt_samples = random.sample(gt_samples, sample_num)
    sample_token_list = [sample['sample_token'] for sample in gt_samples]
    sample_annos_dict = {k: sample_annos_dict[k] for k in sample_token_list}
    # gt annotation visualization
    gt_visualizer(gt_samples, SAVE_PATH, NUSCENES_PATH, SAMPLE_DATA_PATH)
    # load nuscenes info 
    sample_info_dict = nuscenes_info_loader(NUSCENES_PATH, SAMPLE_DATA_PATH)
    # pred annotation visualization
    pred_visualizer(NUSCENES_PATH, SAMPLE_DATA_PATH, SAVE_PATH, sample_annos_dict, sample_info_dict)
    # iou filter
    gt_box_list = get_box_list(gt_samples)
    fp_filter(NUSCENES_PATH, 
              SAMPLE_DATA_PATH, 
              SAVE_PATH, 
              sample_annos_dict, 
              sample_info_dict,
              gt_box_list, 
              iou_threshold=0.3)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_name', type=str, required=True, help="Data path to changan bev")
    parser.add_argument('--exp_name', type=str, required=True, help="exp name")
    parser.add_argument('--sample_num', type=int, required=True, help="sample num you want to visualize")
    parser.add_argument('--eval_name', type=str, required=True, help="eval exp name")
    args = parser.parse_args()
    SAMPLE_DATA_PATH = os.path.join('../data', args.data_name)
    NUSCENES_PATH = os.path.join(SAMPLE_DATA_PATH, 'nuscenes')
    PRED_PATH = os.path.join('../outputs', args.exp_name, "results_nusc.json")
    SAVE_PATH = os.path.join(args.eval_name, "sample_visual")
    sample_num = args.sample_num
    visualize(NUSCENES_PATH, SAMPLE_DATA_PATH, PRED_PATH, SAVE_PATH, sample_num)